/*
  Copyright (C) 2011, Kyungjoo Kim
  
  This file is part of UHM (Unassembled HyperMatrix) sparse direct solver.
  
  All rights reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.
  
  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
  
  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  Also add information on how to contact you by electronic and paper mail.
  
  
  Kyungjoo Kim
  iamkyungjoo@gmail.com
*/
#include "uhm/common.hxx"
#include "uhm/const.hxx"
#include "uhm/util.hxx"

#include "uhm/object.hxx"

#include "uhm/operation/scheduler.hxx"
#include "uhm/operation/mesh.hxx"
#include "uhm/operation/element.hxx"

#include "uhm/mesh/node.hxx"
#include "uhm/mesh/element.hxx"

#include "uhm/matrix/uhm/matrix.hxx"

#include "uhm/mesh/mesh.hxx"
extern "C" {
#include "defs.h"
#include "struct.h"
#include "macros.h"
#include "rename.h"

  extern void SetUpGraph(GraphType *, int, int, int,
                         idxtype *, idxtype *, idxtype *, idxtype *, int);
  extern int idxsum(int, idxtype *);
  extern void InitRandom(int);
  extern void AllocateWorkSpace(CtrlType *, GraphType *, int);
  extern void FreeWorkSpace(CtrlType *, GraphType *);
  extern void MlevelEdgeBisection(CtrlType *, GraphType *, int *, float);
  extern void SplitGraphPart(CtrlType *,
                             GraphType *, GraphType *, GraphType *);
  extern void GKfree(void **,...);

  extern void METIS_PartGraphRecursive(int *, int *, int *, int *,
                                       int *, int *, int *, int *,
                                       int *, int *, int *);

}

#define UHM_PARTITION_MESH_MEASURE_TIME 1

namespace uhm {
  // --------------------------------------------------------------
  // ** Pre-condition
  // 1. Mesh has connectivity information of nodes nad elements.
  // 2. Mssh does NOT build tree yet.

  // ** Post-condition
  // 1. Lock the mesh - it will arrange all connectivity
  bool partition_mesh_var_2(Mesh m, int nparts) {

#ifdef UHM_PARTITION_MESH_MEASURE_TIME 
    printf(":: begin :: partition_mesh_var_1\n");
    double t, time;
#endif

    assert(m && mesh_valid(m));

    // ---------------------------------------------------------------
    // collect leaves and clean the mesh if it contains non-leaf elements

#ifdef UHM_PARTITION_MESH_MEASURE_TIME 
    t = timer();
#endif

    m->remove_non_leaves();
    std::vector< Element > *leaves = new std::vector< Element >;
    leaves->reserve( m->get_n_elements() );
    
    std::map< int, Element_ >::iterator mit;
    for (mit=m->elements.begin();mit!=m->elements.end();++mit) {
      Element elt = &(mit->second);
      leaves->push_back( elt );
      elt->set_generation(0);
      elt->restore_connectivity();
      elt->set_marker(0, -1);
      elt->set_marker(1, -1);
      elt->separate_nodes();
    }

    
#ifdef UHM_PARTITION_MESH_MEASURE_TIME 
    time = timer() - t;
    printf(":: adding leaves %lf\n", time);
#endif
    
    // ---------------------------------------------------------------
    // ** partition into nparts
    if (leaves->size() > 1) {
      int n = leaves->size(), wgtflag = 1, numflag = 0,
        options = 0, edgecut = 0;

      // form a family
      Element parent = m->add_element();
      parent->set_generation(-2);

      for (int i=0;i<nparts;++i) {
        Element elt = m->add_element();
        elt->set_generation(-1);
        elt->set_parent( parent );
        parent->add_child( elt );
      }
      
      // graph
      std::vector< int > xadj, adjncy, adjwgt, part;
      part.assign( n, 0 );

#ifdef UHM_PARTITION_MESH_MEASURE_TIME 
      t = timer();
#endif
      
      build_graph(leaves, parent,
                  &xadj, &adjncy, &adjwgt);
      
#ifdef UHM_PARTITION_MESH_MEASURE_TIME 
      time = timer() - t;
      printf(":: building graph leaves %lf\n", time);
#endif

      // partition
#ifdef UHM_PARTITION_MESH_MEASURE_TIME 
      t = timer();
#endif
      
      METIS_PartGraphRecursive(&n, &xadj[0], &adjncy[0], NULL, &adjwgt[0],
                               &wgtflag, &numflag,
                               &nparts, &options, &edgecut, &part[0]);
      
#ifdef UHM_PARTITION_MESH_MEASURE_TIME 
      time = timer() - t;
      printf(":: METIS graph partitioning %lf\n", time);
#endif

      // each partition hold leaves living in his region
      for (int i=0;i<part.size();++i) {
        Element c = parent->get_child(part.at(i));
        Element l = leaves->at(i);

        c->add_child( l );
        l->set_parent( c );
      }
    }

    delete leaves;


#ifdef UHM_PARTITION_MESH_MEASURE_TIME
    printf(":: end :: partition_mesh_var_1\n");
#endif

    return true;
  }
}
